{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<xarray.DataArray 'T2' (Time: 4, bottom_top: 1, south_north: 96, west_east: 96)>\n",
      "array([[[[303.6255 , ..., 303.81546],\n",
      "         ...,\n",
      "         [304.7401 , ..., 300.15717]]],\n",
      "\n",
      "\n",
      "       ...,\n",
      "\n",
      "\n",
      "       [[[301.89346, ..., 302.56534],\n",
      "         ...,\n",
      "         [302.61246, ..., 298.01028]]]], dtype=float32)\n",
      "Coordinates:\n",
      "  * Time     (Time) datetime64[ns] 2005-08-28 2005-08-28T03:00:00 ...\n",
      "    XTIME    (Time) float64 0.0 180.0 360.0 540.0\n",
      "    XLAT     (south_north, west_east) float32 21.302004 21.302004 21.302004 ...\n",
      "    XLONG    (south_north, west_east) float32 -90.57406 -90.484116 ...\n",
      "Dimensions without coordinates: bottom_top, south_north, west_east\n",
      "Attributes:\n",
      "    FieldType:    104\n",
      "    MemoryOrder:  XY \n",
      "    description:  TEMP at 2 M\n",
      "    units:        K\n",
      "    stagger:      \n",
      "    coordinates:  XLONG XLAT XTIME\n",
      "    projection:   Mercator(stand_lon=-89.0, moad_cen_lat=27.99999237060547, t...\n"
     ]
    }
   ],
   "source": [
    "from __future__ import print_function, division\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "from netCDF4 import Dataset\n",
    "from wrf import getvar, ALL_TIMES, to_np\n",
    "import xarray\n",
    "\n",
    "path = '/scratch/mawagner/2015_GRELL3D/TESTFILES/TestTime'\n",
    "filename_list = os.listdir(path)\n",
    "filename_list.sort()\n",
    "\n",
    "#filename_list = [\"/Users/ladwig/Documents/wrf_files/wrf_vortex_single/wrfout_d02_2005-08-28_00:00:00\",\n",
    "#                \"/Users/ladwig/Documents/wrf_files/wrf_vortex_single/wrfout_d02_2005-08-28_03:00:00\",\n",
    "#                \"/Users/ladwig/Documents/wrf_files/wrf_vortex_single/wrfout_d02_2005-08-28_06:00:00\",\n",
    "#                \"/Users/ladwig/Documents/wrf_files/wrf_vortex_single/wrfout_d02_2005-08-28_09:00:00\"]\n",
    "\n",
    "# Result shape \n",
    "result_shape = (6, 1, 290, 265)\n",
    "#result_shape = (4, 1, 96, 96)\n",
    "\n",
    "# Let's get the first time so we can copy the metadata later\n",
    "f = Dataset(filename_list[0])\n",
    "# By setting squeeze to False, you'll get all the dimension names.\n",
    "z1 = getvar(f, \"T2\", 0, squeeze=False)\n",
    "xlat = getvar(f, \"XLAT\", 0)\n",
    "xlong = getvar(f, \"XLONG\", 0)\n",
    "\n",
    "\n",
    "z_final = np.empty(result_shape, np.float32)\n",
    "\n",
    "# Modify this number if using more than 1 time per file\n",
    "times_per_file = 1\n",
    "#times_per_file = 4\n",
    "\n",
    "data_times = []\n",
    "xtimes = []\n",
    "for timeidx in range(result_shape[0]):\n",
    "    # Compute the file index and the time index inside the file\n",
    "    fileidx = timeidx // times_per_file\n",
    "    file_timeidx = timeidx % times_per_file\n",
    "\n",
    "    f = Dataset(filename_list[fileidx])\n",
    "    z = getvar(f, \"T2\", file_timeidx)\n",
    "    t = getvar(f, \"Times\", file_timeidx)\n",
    "    xt = getvar(f, \"xtimes\", file_timeidx)\n",
    "    data_times.append(to_np(t))\n",
    "    xtimes.append(to_np(xt))\n",
    "    z_final[timeidx,:] = z[:]\n",
    "    f.close()\n",
    "    \n",
    "# Let's make the metadata. Dimension names should copy easily if you set sqeeze to False, \n",
    "# otherwise you can just set them yourself is a tuple of dimension names. Since you wanted\n",
    "# to keep the bottom_top dimension for this 2D variable (which is normally removed), \n",
    "# I'm doing this manually.\n",
    "z_dims = [\"Time\", \"bottom_top\", \"south_north\", \"west_east\"]\n",
    "\n",
    "# Xarray doesn't copy coordinates easily (it always complains about shape mismatches), so do this\n",
    "# manually\n",
    "z_coords = {}\n",
    "z_coords[\"Time\"] = data_times\n",
    "z_coords[\"XTIME\"] = (\"Time\",), xtimes\n",
    "z_coords[\"XLAT\"] = (\"south_north\", \"west_east\"), xlat\n",
    "z_coords[\"XLONG\"] = (\"south_north\", \"west_east\"), xlong\n",
    "z_name = \"T2\"\n",
    "\n",
    "# Attributes copy nicely\n",
    "z_attrs = {}\n",
    "z_attrs.update(z1.attrs)\n",
    "\n",
    "z_with_meta = xarray.DataArray(z_final, coords=z_coords, dims=z_dims, attrs=z_attrs, name=z_name)\n",
    "\n",
    "print(z_with_meta)\n",
    "    \n",
    "    "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
